package xyz.imeans.w.service;

import com.xiaoleilu.hutool.json.JSONArray;
import com.xiaoleilu.hutool.json.JSONObject;
import com.xiaoleilu.hutool.util.StrUtil;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.codec.binary.Base64;
import org.apache.http.HttpHeaders;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.stereotype.Service;
import xyz.imeans.w.annotion.Table;
import xyz.imeans.w.config.CreateDataConfig;
import xyz.imeans.w.config.CreateEnum;
import xyz.imeans.w.config.InitCreateFactory;
import xyz.imeans.w.config.ObjCreateFactory;
import xyz.imeans.w.util.OkHttpCli;

import java.nio.charset.StandardCharsets;
import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.ResultSet;
import java.sql.Statement;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.StringJoiner;
import java.util.UUID;

@Slf4j
@Service
public class DorisService implements CreateService, InitializingBean {

    private final CreateDataConfig createDataConfig;
    private final OkHttpCli okHttpCli;
    public DorisService(CreateDataConfig createDataConfig, OkHttpCli okHttpCli) {
        this.createDataConfig = createDataConfig;
        this.okHttpCli = okHttpCli;
    }

    @Override
    public void afterPropertiesSet() throws Exception {
        InitCreateFactory.initCreateService(CreateEnum.DORIS, this);
    }

    @Override
    public void exec(CreateDataConfig.Config config) {
        int count = config.getCount();
        if(count <= 0) {
            return;
        }
        Connection conn = null;
        Statement stmt = null;
        ResultSet resultSet = null;
        try {
            Class<?> aClass = Class.forName(config.getModel());
            Table table = aClass.getAnnotation(Table.class);
            if(StrUtil.isBlank(table.schma())) {
                return;
            }

            Map<String, String> map = createDataConfig.getSource().get("doris");
            Class.forName(map.get("driver-class-name"));
            conn = DriverManager.getConnection(map.get("url"), map.get("username"), map.get("password"));
            stmt = conn.createStatement();
            // 检查是否存在表
            resultSet = stmt.executeQuery("SELECT COUNT(*) cnt FROM information_schema.tables " +
                    "WHERE table_name = '" + table.value() + "';");
            int cnt = 0;
            if(resultSet.next()) {
                cnt = resultSet.getInt(1);
            }

            // 获取表字段
            Map<String, String> columns = ObjCreateFactory.getFieldColumn(aClass);
            if(columns.size() == 0) {
                return;
            }

            // 创建表
            if(cnt == 0) {
                StringBuilder sb = new StringBuilder("CREATE TABLE ");
                sb.append(table.value()).append(" (\n");
                String uniqueKey = columns.remove("my_system_data_uniqueKey");
                String duplicateKey = columns.remove("my_system_data_duplicateKey");
                String aggregateKey = columns.remove("my_system_data_aggregateKey");
                for (Map.Entry<String, String> entry : columns.entrySet()) {
                    sb.append(entry.getValue());
                    if(StrUtil.isNotBlank(uniqueKey) && entry.getValue().equals(uniqueKey)) {
                        sb.append(" VARCHAR(255) NOT NULL, \n");
                    } else {
                        sb.append(" VARCHAR(255) NULL, \n");
                    }
                }
                sb.append(")  ENGINE=OLAP \n");
                if(StrUtil.isNotBlank(uniqueKey)) {
                    sb.append("UNIQUE KEY(").append(uniqueKey).append(")\n");
                    sb.append("DISTRIBUTED BY HASH(").append(uniqueKey).append(") BUCKETS 2\n");
                } else if(StrUtil.isNotBlank(duplicateKey)) {
                    sb.append("DUPLICATE KEY(").append(duplicateKey).append(")\n");
                    sb.append("DISTRIBUTED BY HASH(").append(duplicateKey).append(") BUCKETS 2\n");
                } else if(StrUtil.isNotBlank(aggregateKey)) {
                    sb.append("AGGREGATE KEY(").append(aggregateKey).append(")\n");
                    sb.append("DISTRIBUTED BY HASH(").append(aggregateKey).append(") BUCKETS 2\n");
                }
                sb.append("PROPERTIES\n");
                sb.append("(\n");
                sb.append(" \"replication_num\" = \"1\"\n");
                sb.append(");");
                System.out.println(sb);
                stmt.execute(sb.toString());
            }

            // 生成数据
            List<Map<String, Object>> list = new ArrayList<>(count);
            for (int i = 0; i < count; i++) {
                Map<String, Object> d = ObjCreateFactory.createObj(aClass);
                if(d.size() > 0) {
                    list.add(d);
                }
            }
            if(list.size() == 0) {
                return;
            }

            // 写入
            String url = String.format("http://%s/api/%s/%s/_stream_load", map.get("host"), table.schma(), table.value());

            JSONArray datas = new JSONArray();
            for (Map<String, Object> objectMap : list) {
                JSONObject json = new JSONObject();
                for (Map.Entry<String, String> column : columns.entrySet()) {
                    json.put(column.getValue(), objectMap.get(column.getKey()));
                }
                datas.add(json);
            }

            Map<String, String> headers = new HashMap<>();
            headers.put("label", UUID.randomUUID().toString());
            StringJoiner sj = new StringJoiner(",");
            for (Map.Entry<String, String> column : columns.entrySet()) {
                sj.add(column.getValue());
            }
            headers.put("columns", sj.toString());
            headers.put("column_separator", ",");
            headers.put("Expect", "100-continue");
            headers.put("format", "json");
            headers.put("strip_outer_array", "true");
            headers.put("num_as_string", "true");
            byte[] bytes = Base64.encodeBase64(String.format("%s:%s", map.get("username"), map.get("password"))
                    .getBytes(StandardCharsets.UTF_8));
            headers.put(HttpHeaders.AUTHORIZATION, "Basic "+new String(bytes));

            String s = okHttpCli.doPutJson(url, datas.toString(), headers);
            log.info("导入结果：{}", s);
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            try {
                if (resultSet != null) {
                    resultSet.close();
                }
                if (stmt != null) {
                    stmt.close();
                }
                if (conn != null) {
                    conn.close();
                }
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
    }
}
